package com.atbeijing.bigdata.spark.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, IntegerType, StructField, StructType}
import org.apache.spark.sql.{Row, SparkSession}

object SparkSQL07_UDAF {

    def main(args: Array[String]): Unit = {

        // TODO 创建环境对象
        val sparkConf = new SparkConf().setMaster("local[*]").setAppName("SparkSQL")
        val spark: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
        import spark.implicits._

        // RDD
        val rdd = spark.sparkContext.makeRDD(
            List(
                (1, "zhangsan", 30),
                (2, "lisi", 40),
                (3, "wangwu", 50)
            )
        )

        val df = rdd.toDF("id", "name", "age")

        df.createOrReplaceTempView("user")

        spark.udf.register("ageAvg", new MyAgeAvgUDAF())
        // sql语句中早期版本只支持弱类型的聚合函数
        spark.sql("select ageAvg(age) from user").show


        spark.stop()
    }
    // TODO 自定义用户聚合函数（UDAF）- 弱类型
    // 1. 继承UserDefinedAggregateFunction类
    // 2. 重写方法
    class MyAgeAvgUDAF extends UserDefinedAggregateFunction {
        // 输入数据的结构
        override def inputSchema: StructType = {
            // age(int)
            StructType(
                Array(
                    StructField("age", IntegerType)
                )
            )
        }

        // 缓冲区数据的结构
        override def bufferSchema: StructType = {
            StructType(
                Array(
                    StructField("total", IntegerType),
                    StructField("cnt", IntegerType)
                )
            )
        }

        // 输出数据的类型
        override def dataType: DataType = IntegerType

        // 数据的稳定性
        override def deterministic: Boolean = true

        // 缓冲区的初始化[total,cnt]
        override def initialize(buffer: MutableAggregationBuffer): Unit = {
            //第0个位置
            buffer.update(0, 0)
            //第一个位置
            buffer.update(1, 0)
        }

        // 使用输入的数据更新缓冲区中的数据[total,cnt]
        override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
            val inputAge = input.getInt(0)
            val bufferAge = buffer.getInt(0)
            val bufferCnt = buffer.getInt(1)
            buffer.update(0, inputAge + bufferAge)
            buffer.update(1, bufferCnt + 1)
        }

        // 合并多个缓冲区[total,cnt]
        override def merge(buffer1: MutableAggregationBuffer, input: Row): Unit = {
            val inputAge = input.getInt(0)
            val inputCnt = input.getInt(1)
            val bufferAge = buffer1.getInt(0)
            val bufferCnt = buffer1.getInt(1)
            buffer1.update(0, inputAge + bufferAge)
            buffer1.update(1, bufferCnt + inputCnt)
        }

        // 计算
        override def evaluate(buffer: Row): Any = {
            buffer.getInt(0) / buffer.getInt(1)
        }
    }
}
